package ai

import (
	"context"
	"encoding/json"
	"fmt"
	"gitee.com/CrazyMage/blog-service/global"
	"github.com/redis/go-redis/v9"
	"time"
)

const (
	CHAT_MEMORY_CONVERSATION_ID_KEY = "chat_memory_conversation_id:"
)

type Message struct {
	Role    string
	Content string
	Time    time.Time
}

type ChatMemory interface {
	//存储单个回话
	AddSingleMessage(conversationId string, message *Message) error
	//存储多个回话
	Add(conversationId string, messages []*Message) error
	//获取回话
	Get(conversationId string, lastN int) ([]*Message, error)
	//清空回话
	Clear(conversationId string) error
}

type RedisChatMemory struct {
	Context     context.Context
	RedisClient *redis.Client
}

func NewRedisChatMemory(ctx context.Context) *RedisChatMemory {
	return &RedisChatMemory{
		Context:     ctx,
		RedisClient: global.RedisClient,
	}
}

func (rc RedisChatMemory) AddSingleMessage(conversationId string, message *Message) error {
	return rc.Add(conversationId, []*Message{message})
}

func (rc RedisChatMemory) Add(conversationId string, messages []*Message) error {

	key := CHAT_MEMORY_CONVERSATION_ID_KEY + conversationId

	// 1. 序列化所有消息
	serialized := make([]interface{}, len(messages))
	for i, msg := range messages {
		data, err := json.Marshal(msg) // 使用JSON序列化
		if err != nil {
			return fmt.Errorf("序列化消息失败: %w", err)
		}
		serialized[i] = data
	}

	// 2. 批量写入Redis
	if err := rc.RedisClient.RPush(rc.Context, key, serialized...).Err(); err != nil {
		return fmt.Errorf("Redis写入失败: %w", err)
	}

	return nil

}

func (rc RedisChatMemory) Get(conversationId string, lastN int) ([]*Message, error) {
	key := CHAT_MEMORY_CONVERSATION_ID_KEY + conversationId

	res, err := rc.RedisClient.LRange(rc.Context, key, -int64(lastN), -1).Result()
	if err != nil {
		return nil, err
	}

	// 反序列化每条消息
	messages := make([]*Message, 0, len(res))
	for _, item := range res {
		var msg Message
		if err = json.Unmarshal([]byte(item), &msg); err != nil {
			return nil, err
		}
		messages = append(messages, &msg)
	}

	return messages, nil
}

func (rc RedisChatMemory) Clear(conversationId string) error {
	key := CHAT_MEMORY_CONVERSATION_ID_KEY + conversationId

	_, err := rc.RedisClient.Del(rc.Context, key).Result()

	return err
}
